from mcp.server.fastmcp import FastMCP
import math
from typing import Union

mcp = FastMCP("Large Number Calculator Service")

@mcp.tool()
def add(a: int, b: int) -> int:
    """大数加法运算
    Args:
        a: 第一个大整数
        b: 第二个大整数
    Returns:
        两数之和
    """
    return a + b

@mcp.tool()
def subtract(a: int, b: int) -> int:
    """大数减法运算
    Args:
        a: 被减数
        b: 减数
    Returns:
        两数之差
    """
    return a - b

@mcp.tool()
def multiply(a: int, b: int) -> int:
    """大数乘法运算
    Args:
        a: 第一个乘数
        b: 第二个乘数
    Returns:
        两数之积
    """
    return a * b

@mcp.tool()
def divide(a: int, b: int) -> Union[int, float]:
    """大数除法运算（整除返回整数，否则返回浮点数）
    Args:
        a: 被除数
        b: 除数（不能为0）
    Returns:
        商（整数或浮点数）
    """
    if b == 0:
        raise ValueError("除数不能为零")
    if a % b == 0:
        return a // b
    return a / b

@mcp.tool()
def power(base: int, exponent: int) -> int:
    """大数乘方运算
    Args:
        base: 底数
        exponent: 指数（非负整数）
    Returns:
        底数的指数次幂
    """
    if exponent < 0:
        raise ValueError("指数必须为非负整数")
    return base ** exponent

@mcp.tool()
def square_root(n: int) -> Union[int, float]:
    """大数开平方运算（返回整数部分或精确浮点数）
    Args:
        n: 被开方数（非负整数）
    Returns:
        平方根（整数部分或精确浮点数）
    """
    if n < 0:
        raise ValueError("被开方数必须为非负整数")
    root = math.isqrt(n)
    if root * root == n:
        return root
    return math.sqrt(n)

@mcp.tool()
def logarithm(n: int, base: int = 10) -> float:
    """大数对数运算（以base为底）
    Args:
        n: 真数（正整数）
        base: 底数（正整数，默认10）
    Returns:
        对数值（浮点数）
    """
    if n <= 0 or base <= 0:
        raise ValueError("真数和底数必须为正整数")
    return math.log(n, base)

@mcp.tool()
def compare(a: int, b: int) -> str:
    """大数比较运算
    Args:
        a: 第一个数
        b: 第二个数
    Returns:
        比较结果（'a > b'/'a == b'/'a < b'）
    """
    if a > b:
        return 'a > b'
    elif a == b:
        return 'a == b'
    else:
        return 'a < b'

if __name__ == "__main__":
    mcp.run(transport="stdio")